# coding=utf-8
import numpy as np


def initialization_centroid(data, k):
    """
    初始化质心
    :param data:
    :param k:
    :return:
    """
    m_data = np.mat(data)
    centroid_coordinate_T = []  # 质心坐标的转置
    for r in range(len(data[0])):
        r_max = np.max(m_data[:, r])  # 取列的最大值
        r_min = np.min(m_data[:, r])  # 取列的最小值
        auto_distant = (r_max - r_min) / k  # 等分长度  每个维度距离被质心个数等分
        auto_coordinate = []
        for n in range(k):
            r_min += auto_distant  # 计算质心坐标
            auto_coordinate.append(r_min)
        centroid_coordinate_T.append(auto_coordinate)  # 质心
    return np.transpose(centroid_coordinate_T).tolist()


def start_clustering(data, centroid_list, k):
    """
    初始聚类
    :param k:
    :param centroid_list:
    :param data:
    :return:
    """

    cluster_dic = {}
    for i in range(k):
        cluster_dic[str(i)] = []  # 根据簇index 初始化聚类簇

    for line in data:
        vector_length = []
        for c in centroid_list:
            # print(c, line)
            length = []
            if c and line:
                length = np.sum(np.square(np.mat(c) - np.mat(line)))  # 计算欧氏距离的平方
            vector_length.append(length)
        index = vector_length.index(max(vector_length))  # 取出位置索引index
        cluster_dic[str(index)].append(line)  # 聚类
    return cluster_dic


def calculate_new_centroid(cluster_dic, k):
    """
    计算新质心
    :param cluster_dic:
    :param k:
    :return:
    """
    new_centroid_coordinate = []  # 质心坐标的转置
    for cl in cluster_dic.values():
        new_auto_coordinate = []
        if cl:
            mcl = np.mat(cl)
            # print(cl)
            for i in range(len(cl[0])):
                new_auto_coordinate.append(np.mean(mcl[:, i]))  # 计算每一行的平均值
        new_centroid_coordinate.append(new_auto_coordinate)
    return new_centroid_coordinate  # 平均值作为新的质心


def iterative_clustering(data, iteration_times, k):
    """
    迭代训练
    :param data:
    :param iteration_times:
    :param k:
    :return:
    """
    centroid = initialization_centroid(data, k)
    cluster_dic = start_clustering(data, centroid, k)  # 初始化聚类，得到初始质心和簇
    # transient_cluster = None
    while iteration_times > 0:
        new_centroid = calculate_new_centroid(cluster_dic, k)  # 迭代获取质心
        cluster_dic = start_clustering(data, new_centroid, k)  # 迭代分簇
        iteration_times -= 1
        # print(iteration_times)
    return cluster_dic


if __name__ == '__main__':
    k_data = [
        [13, 45, 98],
        [23, 77, 42],
        [55, 56, 65],
        [89, 64, 32],
        [21, 11, 13],
        [21, 47, 56],
        [86, 55, 34],
        [11, 11, 46],
        [65, 43, 56]
    ]

    data2 = [
        [6.99438039, 5.05456275],
        [8.08169456, 7.97506735],
        [3.02211698, 6.00770189],
        [2.95832148, 2.98598456],
        [2.9750599, 3.77881139],
        [7.311725, 5.00685],
        [6.7122963, 5.09697407],
        [8.08169456, 7.97506735],
    ]

    cl = iterative_clustering(data2, 2000, 2)
    print('result:', cl)
